"""
A set of helpers to read and save data generated
during a pqsheet sessions.

The basic Idea is to have a structure like this:
- locVars element, the main node;
- a node for each variable in localVars.

In future I will record also the data in
the importer window.
"""
from xml.etree import ElementTree as ET
import numpy as np, string, datetime
from dateutil import parser
import string

from PyQt4 import QtCore

floatTypes = (float,np.float64)
intTypes = (int,np.int32)

#strTypes = (str, QtCore.QString)#BHC
strTypes = (str, str)#BHC

class xmlDataHandler:
    """
    An helper class for reading and writing
    xml data.
    """
    def __init__(self,xmlFileName):
        self.__theFileName = xmlFileName
    
    def __getSingleDirection__(self,value):
        return self.__encodeType__(value), \
            string.join([str(i) for i in value],' ')
    
    def __getMultiDirection__(self,value):
        valtype = self.__encodeType__(value)
        if valtype == 'array':
            revalue = value.reshape(-1)
            return self.__encodeType__(revalue[0]),\
                string.join([str(i) for i in revalue])
        if (len(np.shape(value)) == 2):
            return self.__getDoubleListDirection__(value)
        else:
            raise Exception('Saving not allowed for lists\nwith shape greater than 2.\nYou can convert it in an np.array.')
    
    def __getDoubleListDirection__(self,value):
        out = [self.__getSingleDirection__(i)[1] for i in value]
        elType = self.__getSingleDirection__(value[0])[0]
        return elType, string.join(out,' ')
    
    def __formatText__(self,value,shape):
        """
        Dimensions greater than 2 are not supported for the moment.
        """
        if len(shape) == 1:
            return self.__getSingleDirection__(value)
        elif len(shape) > 1:
            return self.__getMultiDirection__(value)        
        return self.__encodeType__(value), str(value)
    
    def __encodeType__(self,value):
        """
        Maybe there's a better way to extract
        the type string!?
        """
        valType = type(value)
        if valType is np.ndarray: return 'array'
        elif valType is list: return 'list'
        elif valType in floatTypes: return 'float'
        elif valType in intTypes: return 'int'
        elif valType in strTypes: return 'str'#BHC
        elif valType is datetime.date: return 'date'
        return 'unknown'
    
    def __encodeShape__(self,shape):
        if not shape: return '0'
        return string.join([str(i) for i in shape],' ')

    def __decodeScalar__(self,valtype,text):
        if valtype == 'float':
            print "Float decoded"
            return string.atof(text)
        elif valtype == 'int':
            return string.atoi(text)
        elif valtype == 'str':#BHC
            return str(text)#BHC
        elif valtype == 'date':
            return parser.parse(text).date()
    
    def __decodeVector__(self,shape,valtype,eltype,text):
        print "decodeVector called"
        out = []
        elements = text.split()
        for i in elements:
            out.append(self.__decodeScalar__(eltype,i))
        if valtype == 'array': return np.array(out)
        return out

    def __decodeArray__(self,shape,valtype,eltype,text):
        print "decodeArray called"
        if valtype == 'array':
            elements = text.split()
            out = np.array([self.__decodeScalar__(eltype,i) for i in elements])
            return out.reshape(shape)
        else:
            return self.__decodeMultiList__(shape,valtype,eltype,text)

    def __decodeMultiList__(self,shape,eltype,text):
        print "decodeMultiList called"
        if len(shape) > 2:
            raise Exception('Loading not allowed for lists\nwith shape greater than 2.')
        out = []
        elements = text.split()
        for i in range(shape[0]):
            text = string.join(elements[i*shape[1]:(i+1)*shape[1]],' ')
            out.append(self.__decodeVector__(shape[1],type,eltype,text))
        return out
    
    def __decodeContents__(self,attributes,text):
        """
        Reading the variable contents.
        """
        shape = [string.atoi(i) for i in attributes['shape'].split()]
        print "shape: ", shape
        print "i: ", i
        if len(shape) == 1: return self.__decodeVector__(shape[0],attributes['type'],attributes['eltype'],text)
        elif len(shape) > 1: return self.__decodeArray__(shape,attributes['type'],attributes['eltype'],text)
        return self.__decodeScalar__(attributes['type'],text)
    
    def readData(self):
        iFile = open(self.__theFileName)
        tree = ET.parse(iFile)
        root_element = tree.getroot()
        outDict = {}
        for child_element in root_element:
            varName = child_element.attrib['name']
            outDict[varName] = self.__decodeContents__(child_element.attrib,\
                                                     child_element.text)
        iFile.close()
        return outDict
    
    def writeData(self,pubVars):
        oFile = open(self.__theFileName,'w')
        root_element = ET.Element("pubvars")
        for name, value in pubVars.iteritems():
            varShape = np.shape(value)
            elType, outText = self.__formatText__(value,varShape)
            varAttributes = {'name': '%s' % name,
                             'type': '%s' % self.__encodeType__(value),
                             'shape' : '%s' % self.__encodeShape__(varShape),
                             'eltype' : '%s' % elType}
            child_element = ET.Element("var",varAttributes)
            child_element.text = outText
            root_element.append(child_element)
        oFile.write(ET.tostring(root_element))
        oFile.close()